# from gensim.models import Word2Vec
import pickle
import gensim
# import numpy as np
# from numpy import float32 as REAL
import click


@click.command()
@click.option("-p", "--path", type=str, help="The pkl data path")
@click.option("-n", "--name", type=str, help="The output w2v name")
def run(path, name):
    # model = gensim.models.Word2Vec.load("word2vec-train-ida.model")
    # old_key = set(model.wv.index_to_key)

    cs_words = [
        "csqword", "csdword", "cspword", "csxmmword", "csymmword", "csunk", "csoff", "csbyte", "aAssertFail", "csasc"
    ]
    token_list = []
    f = open(path, "rb")
    trace_data = pickle.load(f)
    for file in trace_data:
        funcs = trace_data[file]
        for func in funcs:
            tmp_trace = trace_data[file][func]['blocks']#遍历函数的所有block，并生成语句
            for i in range(len(tmp_trace)):
                for j in range(len(tmp_trace[i])):
                    if tmp_trace[i][j].isdigit() or "0x" in tmp_trace[i][j]:#找数字，用"num代替"
                        tmp_trace[i][j] = "num"
                    for cs_word in cs_words:
                        if cs_word in tmp_trace[i][j]:
                            tmp_trace[i][j] = cs_word
                            break

            token_list.extend(tmp_trace)
    f.close()

    # f = open("tvm_data/save_dir/trace_data_all.pkl", "rb")
    # trace_data = pickle.load(f)
    # for file in trace_data:
    #     funcs = trace_data[file]
    #     for func in funcs:
    #         tmp_trace = trace_data[file][func]['blocks']
    #         token_list.extend(tmp_trace)
    # f.close()

    # model.train(token_list)
    model = gensim.models.Word2Vec(sentences=token_list, vector_size=200, window=10, min_count=1, workers=8)
    # model._save_specials()
    # model.build_vocab(token_list, update=True)
    # length = len(model.wv.index_to_key)
    # model.wv.vectors_lockf = np.zeros(length, dtype=REAL)

    # for i, k in enumerate(model.wv.index_to_key):
    #     if k not in old_key:
    #         model.wv.vectors_lockf[i] = 1.

    # model.train(token_list, total_examples=model.corpus_count, epochs=model.epochs)
    model.save(name)


if __name__ == "__main__":
    run()
